﻿//-------------------------------------------------------------------------
// SimpleRowCount Data Collector -  Integrate data into Load Test Analyzer
// 1. Retrieve the row count from a SQL table before and after a test run
// 2. Write the data to a table in XML format
// 3. Send the xml file back to client 
// 
// Blogs related to the sample code
//  http://blogs.msdn.com/yutong/archive/2010/04/08/integrate-data-collector-data-into-load-test-analyzer-tables-view.aspx
//  http://blogs.msdn.com/yutong/archive/2010/04/08/create-a-simple-data-collector-to-integrate-data-into-load-test-analyzer-tables-view.aspx
//-------------------------------------------------------------------------
using Microsoft.VisualStudio.TestTools.Common;
using Microsoft.VisualStudio.TestTools.Execution;
using System.Collections;
using System.Collections.Generic;
using System.Text;
using System.Xml;
using System.IO;
using System;

namespace RowCountDataCollector
{
    // Provide a URI and friendly name for your diagnostic data adapter
    [DataCollectorTypeUri("datacollector://MyCompany/SimpleRowCountDataCollector/1.0")]
    [DataCollectorFriendlyName("RowCountDataCollector", false)]
    [DataCollectorConfigurationEditor("configurationeditor://Microsoft/GenericEditor/1.0")]
    public class SimpleRowCountDataCollector : DataCollector
    {
        private DataCollectionEvents dataEvents;
        private DataCollectionLogger dataLogger;
        private DataCollectionSink dataSink;
        private XmlElement configurationSettings;
        private int rowCountBeforeTestRun;
        private int rowCountAfterTestRun;
        private string server;
        private string database;
        private string table;

        // Required method called by the testing framework
        public override void Initialize(
            XmlElement configurationElement,
            DataCollectionEvents events,
            DataCollectionSink sink,
            DataCollectionLogger logger,
            DataCollectionEnvironmentContext environmentContext)
        {
            // 1. Init members
            dataEvents = events; // The test events
            dataLogger = logger; // The error and warning log
            dataSink = sink;     // Saves collected data

            // Configuration from the test settings
            configurationSettings = configurationElement;

            // 2. Read all the collector configurations from the Test Settings
            CollectorConfigurationManager configManager = new CollectorConfigurationManager(configurationSettings);
            server = configManager["server"];
            if (string.IsNullOrEmpty(server))
            {
                server = Environment.MachineName;
            }

            database = configManager["database"];
            table = configManager["table"];

            // 3. Register TestCaseStart and TestCaseEnd events for the data collector
            dataEvents.TestCaseStart +=
                new EventHandler<TestCaseStartEventArgs>(OnTestCaseStart);

            dataEvents.TestCaseEnd +=
                new EventHandler<TestCaseEndEventArgs>(OnTestCaseEnd);
        }

        protected override void Dispose(bool disposing)
        {
            if (disposing)
            {
                // unregister events
                dataEvents.TestCaseStart -=
                    new EventHandler<TestCaseStartEventArgs>(OnTestCaseStart);

                dataEvents.TestCaseEnd -=
                    new EventHandler<TestCaseEndEventArgs>(OnTestCaseEnd);
            }
        }

        #region Event Handlers

        // Get the count when the load test is started
        public void OnTestCaseStart(object sender, TestCaseEventArgs e)
        {
            dataLogger.LogWarning(e.Context, "OnSessionStart. Time:" + DateTime.Now.ToString());
            rowCountBeforeTestRun = SQLHelper.GetRowCount(server, database, table);
        }

        // Write the counts to a table in XML format
        public void OnTestCaseEnd(object sender, TestCaseEndEventArgs e)
        {
            // Get the row count after run ends
            rowCountAfterTestRun = SQLHelper.GetRowCount(server, database, table);

            // create the table to be integrated into load test database
            string filePath = Path.Combine(Path.GetTempPath(), "rowCounts.xml");
            RowCountXMLWriter.CreateDataTable();
            RowCountXMLWriter.AddDataRow("BeforeTestStart", rowCountBeforeTestRun);
            RowCountXMLWriter.AddDataRow("AfterTestEnd", rowCountAfterTestRun);
            RowCountXMLWriter.WriteRowCountsToXmlFile(filePath);

            // Send the file back
            dataSink.SendFileAsync(e.Context, filePath, false);
        }
        #endregion
    }
}